{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "93a238e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "84717d3b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TokenizedResult(advantage=-1.0, chat='<|im_start|>system\\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\\n<|im_start|>user\\nWhat is the capital of France?<|im_end|>\\n<|im_start|>assistant\\nLondon', tokens=['<|im_start|>', 'system', '\\n', 'You', ' are', ' Q', 'wen', ',', ' created', ' by', ' Alibaba', ' Cloud', '.', ' You', ' are', ' a', ' helpful', ' assistant', '.', '<|im_end|>', '\\n', '<|im_start|>', 'user', '\\n', 'What', ' is', ' the', ' capital', ' of', ' France', '?', '<|im_end|>', '\\n', '<|im_start|>', 'assistant', '\\n', 'London'], token_ids=[151644, 8948, 198, 2610, 525, 1207, 16948, 11, 3465, 553, 54364, 14817, 13, 1446, 525, 264, 10950, 17847, 13, 151645, 198, 151644, 872, 198, 3838, 374, 279, 6722, 315, 9625, 30, 151645, 198, 151644, 77091, 198, 39572], input_pos=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36], assistant_mask=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1], logprobs=[nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan], pixel_values=None, image_grid_thw=None, weight=1.0, prompt_id=0, prompt_length=35)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TokenizedResult(advantage=1.0, chat='<|im_start|>system\\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\\n<|im_start|>user\\nWhat is the capital of France?<|im_end|>\\n<|im_start|>assistant\\nParis', tokens=['<|im_start|>', 'system', '\\n', 'You', ' are', ' Q', 'wen', ',', ' created', ' by', ' Alibaba', ' Cloud', '.', ' You', ' are', ' a', ' helpful', ' assistant', '.', '<|im_end|>', '\\n', '<|im_start|>', 'user', '\\n', 'What', ' is', ' the', ' capital', ' of', ' France', '?', '<|im_end|>', '\\n', '<|im_start|>', 'assistant', '\\n', 'Paris'], token_ids=[151644, 8948, 198, 2610, 525, 1207, 16948, 11, 3465, 553, 54364, 14817, 13, 1446, 525, 264, 10950, 17847, 13, 151645, 198, 151644, 872, 198, 3838, 374, 279, 6722, 315, 9625, 30, 151645, 198, 151644, 77091, 198, 59604], input_pos=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36], assistant_mask=[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1], logprobs=[nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, -0.01], pixel_values=None, image_grid_thw=None, weight=1.0, prompt_id=0, prompt_length=35)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from IPython.display import clear_output\n",
    "from openai.types.chat.chat_completion import Choice, ChoiceLogprobs\n",
    "from openai.types.chat.chat_completion_message import ChatCompletionMessage\n",
    "from openai.types.chat.chat_completion_token_logprob import ChatCompletionTokenLogprob\n",
    "from transformers import AutoTokenizer\n",
    "from transformers.utils.logging import disable_progress_bar\n",
    "\n",
    "import art\n",
    "from art.preprocessing.tokenize import tokenize_trajectory_groups\n",
    "\n",
    "disable_progress_bar()\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen2.5-7B-Instruct\")\n",
    "\n",
    "tokenized_results = list(\n",
    "    tokenize_trajectory_groups(\n",
    "        tokenizer,\n",
    "        [\n",
    "            art.TrajectoryGroup(\n",
    "                [\n",
    "                    art.Trajectory(\n",
    "                        messages_and_choices=[\n",
    "                            {\n",
    "                                \"role\": \"user\",\n",
    "                                \"content\": \"What is the capital of France?\",\n",
    "                            },\n",
    "                            {\"role\": \"assistant\", \"content\": \"London\"},\n",
    "                        ],\n",
    "                        reward=0.0,\n",
    "                    ),\n",
    "                    art.Trajectory(\n",
    "                        messages_and_choices=[\n",
    "                            {\n",
    "                                \"role\": \"user\",\n",
    "                                \"content\": \"What is the capital of France?\",\n",
    "                            },\n",
    "                            Choice(\n",
    "                                finish_reason=\"stop\",\n",
    "                                index=0,\n",
    "                                logprobs=ChoiceLogprobs(\n",
    "                                    content=[\n",
    "                                        ChatCompletionTokenLogprob(\n",
    "                                            token=\"token:59604\",\n",
    "                                            bytes=[80, 97, 114, 105, 115],\n",
    "                                            logprob=-0.01,\n",
    "                                            top_logprobs=[],\n",
    "                                        )\n",
    "                                    ]\n",
    "                                ),\n",
    "                                message=ChatCompletionMessage(\n",
    "                                    content=\"Paris\",\n",
    "                                    role=\"assistant\",\n",
    "                                ),\n",
    "                            ),\n",
    "                        ],\n",
    "                        reward=1.0,\n",
    "                    ),\n",
    "                ]\n",
    "            )\n",
    "        ],\n",
    "        allow_training_without_logprobs=True,\n",
    "        scale_rewards=True,\n",
    "        shuffle_group_trajectories=False,\n",
    "    )\n",
    ")\n",
    "for result in tokenized_results:\n",
    "    result.advantage = round(result.advantage, 2)\n",
    "    result.weight = round(result.weight, 2)\n",
    "    # set prompt_id to 0 to eliminate stochasticity\n",
    "    result.prompt_id = 0\n",
    "    display(result)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
